'''
Author: duliang thinktanker@163.com
Date: 2025-10-07 19:20:17
LastEditors: duliang thinktanker@163.com
LastEditTime: 2025-10-07 21:40:43
FilePath: \noise-analysis\record.py
Description: 这是默认设置,请设置`customMade`, 打开koroFileHeader查看配置 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE
'''

import pyaudio
import wave
import numpy as np
import sys


class SpeakerRecorder:

    def __init__(self,
                 chunk=1024,
                 format=pyaudio.paInt16,
                 channels=2,
                 rate=44100):
        self.chunk = chunk
        self.format = format
        self.channels = channels
        self.rate = rate
        self.audio = pyaudio.PyAudio()

    def find_loopback_device(self):
        """
        查找系统的环回设备（立体声混音、虚拟音频设备等）
        返回设备索引，如果找不到则返回None
        """
        loopback_keywords = [
            'stereo mix', '立体声混音', 'loopback', 'what u hear', 'what you hear',
            'monitor', 'virtual cable', 'voicemeeter'
        ]

        # 首先查找完全匹配的设备
        for i in range(self.audio.get_device_count()):
            info = self.audio.get_device_info_by_index(i)
            device_name = info['name'].lower()

            # 检查是否包含环回关键词并且是输入设备
            if (any(keyword in device_name for keyword in loopback_keywords)
                    and info['maxInputChannels'] > 0):
                # 检查设备是否支持指定的采样率
                try:
                    if info['defaultSampleRate'] == float(self.rate):
                        print(f"找到匹配的环回设备: {info['name']} (索引: {i})")
                        return i
                except KeyError:
                    # 如果没有defaultSampleRate属性，则直接返回
                    print(f"找到可能的环回设备: {info['name']} (索引: {i})")
                    return i

        # 如果没找到关键词匹配的设备，则尝试查找第一个可用的输入设备（除了麦克风）
        for i in range(self.audio.get_device_count()):
            info = self.audio.get_device_info_by_index(i)
            if (info['maxInputChannels'] > 0
                    and 'microphone' not in info['name'].lower()
                    and info['defaultSampleRate'] == float(self.rate)):
                print(f"使用备选输入设备: {info['name']} (索引: {i})")
                return i

        # 最后，如果没有找到完全匹配的设备，返回第一个环回设备（即使采样率不匹配）
        for i in range(self.audio.get_device_count()):
            info = self.audio.get_device_info_by_index(i)
            device_name = info['name'].lower()

            if (any(keyword in device_name for keyword in loopback_keywords)
                    and info['maxInputChannels'] > 0):
                print(f"警告: 使用采样率不匹配的环回设备: {info['name']} (索引: {i})")
                print(
                    f"      设备采样率: {info.get('defaultSampleRate', 'Unknown')}, 请求采样率: {self.rate}"
                )
                return i

        return None

    def record_speaker_output(self, filename, record_seconds):
        """
        录制扬声器输出音频
        注意: 这需要系统支持环回录音或者虚拟音频设备
        """
        # 查找合适的输入设备
        input_device_index = self.find_loopback_device()

        if input_device_index is None:
            print("错误: 未找到合适的音频输入设备")
            print("请检查是否启用了'立体声混音'或其他环回设备")
            return False

        # 获取设备信息用于调试
        device_info = self.audio.get_device_info_by_index(5)
        print(f"使用设备信息: {device_info}")

        stream = None
        try:
            # 在Windows上通常需要指定特定的输入设备索引作为"立体声混音"
            # 在Linux上可以使用pulse audio的monitor source
            # 在macOS上比较复杂，通常需要soundflower或类似工具

            stream = self.audio.open(
                format=self.format,
                channels=self.channels,
                rate=self.rate,
                input=True,
                input_device_index=input_device_index,  # 添加这行来指定设备
                frames_per_buffer=self.chunk)

            print("开始录制扬声器音频...")
            frames = []

            for i in range(0, int(self.rate / self.chunk * record_seconds)):
                data = stream.read(self.chunk)
                frames.append(data)

            print("录制完成!")

        except Exception as e:
            print(f"录制过程中发生错误: {e}")
            print("设备信息:")
            print(f"  设备索引: {input_device_index}")
            print(f"  设备名称: {device_info.get('name', 'Unknown')}")
            print(
                f"  最大输入通道: {device_info.get('maxInputChannels', 'Unknown')}")
            print(
                f"  默认采样率: {device_info.get('defaultSampleRate', 'Unknown')}")
            if stream:
                try:
                    stream.stop_stream()
                    stream.close()
                except:
                    pass
            return False
        finally:
            if stream:
                try:
                    stream.stop_stream()
                    stream.close()
                except Exception as e:
                    print(f"关闭音频流时发生错误: {e}")

        # 保存为WAV文件
        try:
            wf = wave.open(filename, 'wb')
            wf.setnchannels(self.channels)
            wf.setsampwidth(self.audio.get_sample_size(self.format))
            wf.setframerate(self.rate)
            wf.writeframes(b''.join(frames))
            wf.close()
            print(f"音频已保存至 {filename}")
            return True
        except Exception as e:
            print(f"保存音频文件时发生错误: {e}")
            return False

    def list_audio_devices(self):
        """列出所有音频设备"""
        print("可用的音频设备:")
        for i in range(self.audio.get_device_count()):
            info = self.audio.get_device_info_by_index(i)
            print(
                f"{i}: {info['name']} - 输入通道: {info['maxInputChannels']}, 输出通道: {info['maxOutputChannels']}"
            )

        # 特别提示可能的环回设备
        print("\n推荐用于录制系统音频的设备:")
        loopback_device = self.find_loopback_device()
        if loopback_device is not None:
            info = self.audio.get_device_info_by_index(loopback_device)
            print(f"  -> {loopback_device}: {info['name']}")
        else:
            print("  未找到合适的环回设备，请确认系统音频设置")


def main():
    recorder = SpeakerRecorder()

    # 显示可用设备
    recorder.list_audio_devices()

    # 录制音频
    filename = "speaker_output.wav"
    record_time = 10  # 录制10秒

    try:
        success = recorder.record_speaker_output(filename, record_time)
        if not success:
            print("录制失败，请检查系统音频设置")
            print("在Windows上，请确保已启用'立体声混音'设备")
            print("在其他系统上，可能需要安装虚拟音频设备软件")
    except Exception as e:
        print(f"录制过程中发生错误: {e}")
        print("请确保您的系统支持环回录音，或安装虚拟音频设备")


if __name__ == "__main__":
    main()
